{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "206a7709-e8e2-4aec-b755-d72ea1766a48",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import torch\n",
    "from torch import nn\n",
    "from torch.utils.data import DataLoader\n",
    "from torchvision import datasets\n",
    "from torchvision.transforms import ToTensor"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d7e9e123-5791-4807-a4a4-b2f0efa251e0",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 1. 多库含数据集, 教程使用 torchvision.datasets\n",
    "\n",
    "# Download training data from open datasets.\n",
    "training_data = datasets.FashionMNIST(\n",
    "    root=\"data\",\n",
    "    train=True,\n",
    "    download=True,\n",
    "    transform=ToTensor(),\n",
    ")\n",
    "\n",
    "# not Download\n",
    "test_data = datasets.FashionMNIST(\n",
    "    root=\"data\",\n",
    "    train=False,\n",
    "    download=False,\n",
    "    transform=ToTensor(),\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3d72e462-29da-4cb0-88e0-1206462a1022",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# 2. torch.utils.data.DataLoader是torchvision.datasets 的封装\n",
    "\n",
    "batch_size = 64\n",
    "\n",
    "# Create data loaders.\n",
    "train_dataloader = DataLoader(training_data, batch_size=batch_size)\n",
    "test_dataloader = DataLoader(test_data, batch_size=batch_size)\n",
    "\n",
    "for X, y in test_dataloader:\n",
    "    print(f\"Shape of X [N, C, H, W]: {X.shape}\")\n",
    "    print(f\"Shape of y: {y.shape} {y.dtype}\")\n",
    "    break"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "906e3073-b1b0-4a27-87e4-fc15f2772b57",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# 3. 选择加速设备\n",
    "# accelerator 在torch=2.1.0中未提供, 可升级2.6.0\n",
    "device = torch.accelerator.current_accelerator().type if torch.accelerator.is_available() else \"cpu\"\n",
    "print(f\"Using {device} device\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e11f5cf1-3ee0-4dda-bacc-c93573d050df",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 4. Define model, 包含__init__()中定义神经网络层, forward定义张量在模型的各层之间如何向前传播\n",
    "class NeuralNetwork(nn.Module):\n",
    "    def __init__(self):\n",
    "        super().__init__()\n",
    "        self.flatten = nn.Flatten()\n",
    "        self.linear_relu_stack = nn.Sequential(\n",
    "            nn.Linear(28*28, 512),\n",
    "            nn.ReLU(),\n",
    "            nn.Linear(512, 512),\n",
    "            nn.ReLU(),\n",
    "            nn.Linear(512, 10)\n",
    "        )\n",
    "\n",
    "    def forward(self, x):\n",
    "        x = self.flatten(x)\n",
    "        logits = self.linear_relu_stack(x)\n",
    "        return logits\n",
    "\n",
    "model = NeuralNetwork().to(device)\n",
    "print(model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5b631225-b2e1-469a-a5e7-6122be51a3fa",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# 5. 定义损失函数和参数优化器\n",
    "loss_fn = nn.CrossEntropyLoss() # 交叉熵损失函数 \n",
    "optimizer = torch.optim.SGD(model.parameters(), lr=0.001) # stochastic gradient descent 随机梯度下降"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "aa26fc35-9574-4ee4-be0f-3b9138355c17",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# 6. 定义训练过程\n",
    "def train(dataloader, model, loss_fn, optimizer):\n",
    "    size = len(dataloader.dataset)\n",
    "    model.train()  # 设置模型用于训练\n",
    "    for batch, (X, y) in enumerate(dataloader):\n",
    "        X, y = X.to(device), y.to(device)  # 移动张量到设备\n",
    "\n",
    "        # Compute prediction error 计算模型的预测误差\n",
    "        pred = model(X)\n",
    "        loss = loss_fn(pred, y)\n",
    "\n",
    "        # Backpropagation\n",
    "        loss.backward()\n",
    "        optimizer.step()  # 执行优化\n",
    "        optimizer.zero_grad()  # 重置优化器梯度\n",
    "\n",
    "        if batch % 500 == 0:\n",
    "            loss, current = loss.item(), (batch + 1) * len(X)  # 获取损失值loss.item()\n",
    "            print(f\"loss: {loss:>7f}  [{current:>5d}/{size:>5d}]\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f4be6a27-af24-4fed-96fa-7a1ee428ecc4",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# 7. 检查训练结果\n",
    "\n",
    "def test(dataloader, model, loss_fn):\n",
    "    size = len(dataloader.dataset)\n",
    "    num_batches = len(dataloader)\n",
    "    model.eval()   # 设置模型用于计算\n",
    "    test_loss, correct = 0, 0\n",
    "    with torch.no_grad():  # 不再使用反向推导, 计算梯度\n",
    "        for X, y in dataloader:\n",
    "            X, y = X.to(device), y.to(device)\n",
    "            pred = model(X)\n",
    "            test_loss += loss_fn(pred, y).item()\n",
    "            correct += (pred.argmax(1) == y).type(torch.float).sum().item()\n",
    "    test_loss /= num_batches\n",
    "    correct /= size\n",
    "    print(f\"Test Error: \\n Accuracy: {(100*correct):>0.1f}%, Avg loss: {test_loss:>8f} \\n\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "86341aea-dedf-4cca-b130-d92aecbea62a",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# 8. 训练和验证\n",
    "epochs = 5\n",
    "for t in range(epochs):\n",
    "    print(f\"Epoch {t+1}\\n-------------------------------\")\n",
    "    train(train_dataloader, model, loss_fn, optimizer)\n",
    "    test(test_dataloader, model, loss_fn)\n",
    "print(\"Done!\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4594dcae-5c50-435f-adbe-854829d14a09",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# 9. 保存模型\n",
    "\n",
    "torch.save(model.state_dict(), \"model.pth\")\n",
    "print(\"Saved PyTorch Model State to model.pth\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "332f2a65-a45d-4c42-aee2-7e758837b442",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# 10. 加载模型\n",
    "\n",
    "model = NeuralNetwork().to(device)\n",
    "model.load_state_dict(torch.load(\"model.pth\", weights_only=True))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "14b5f9e3-270d-48e2-933b-a2c04c2a16cc",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# 11. 推理\n",
    "\n",
    "classes = [\n",
    "    \"T-shirt/top\",\n",
    "    \"Trouser\",\n",
    "    \"Pullover\",\n",
    "    \"Dress\",\n",
    "    \"Coat\",\n",
    "    \"Sandal\",\n",
    "    \"Shirt\",\n",
    "    \"Sneaker\",\n",
    "    \"Bag\",\n",
    "    \"Ankle boot\",\n",
    "]\n",
    "\n",
    "model.eval()\n",
    "x, y = test_data[0][0], test_data[0][1]\n",
    "with torch.no_grad():\n",
    "    x = x.to(device)\n",
    "    pred = model(x)\n",
    "    predicted, actual = classes[pred[0].argmax(0)], classes[y]\n",
    "    print(f'Predicted: \"{predicted}\", Actual: \"{actual}\"')"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "py312",
   "language": "python",
   "name": "py312"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
